import os
import sys
import base64
from io import BytesIO

import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler  # 引入调度器，优化采样
from transformers import ChineseCLIPProcessor, ChineseCLIPTextModel


# 1. 环境与版本检查
print(f"Python版本: {sys.version}")
print(f"PyTorch版本: {torch.__version__}")
try:
    import diffusers
    print(f"diffusers版本: {diffusers.__version__}")
except ImportError:
    print("diffusers未安装")
try:
    import transformers
    print(f"transformers版本: {transformers.__version__}")
except ImportError:
    print("transformers未安装")

# 检查设备（CPU环境）
device = "cpu"
print(f"使用设备: {device}")
print("提示：CPU生成高质量图像会较慢（可能需要10-30分钟），请耐心等待...\n")

model_path = "./utils/local_server/models"
def check_model_integrity():
    '''模型文件完整性校验'''
    required_files = [
        os.path.join(model_path, "text_encoder", "pytorch_model.bin"),
        os.path.join(model_path, "unet", "diffusion_pytorch_model.bin"),
        os.path.join(model_path, "vae", "diffusion_pytorch_model.bin"),
        os.path.join(model_path, "tokenizer", "vocab.txt")
    ]
    missing = [f for f in required_files if not os.path.exists(f)]
    if missing:
        print("\n❌ 缺失关键模型文件：")
        for f in missing:
            print(f"- {f}")
        print("\n解决方法：重新克隆完整模型")
        print("git clone https://huggingface.co/iic/multi-modal_chinese_stable_diffusion_v1.0 ./models")
        sys.exit(1)
    print("✅ 模型文件完整性检查通过\n")

check_model_integrity()


# 3. 库版本兼容性检查
required_versions = {
    "diffusers": "0.25.0",
    "transformers": "4.36.2",
    "huggingface_hub": "0.19.4"
}
try:
    import diffusers, transformers, huggingface_hub
    if (diffusers.__version__ != required_versions["diffusers"] or
        transformers.__version__ != required_versions["transformers"] or
        huggingface_hub.__version__ != required_versions["huggingface_hub"]):
        print(f"⚠️ 库版本不兼容，建议安装指定版本：")
        print(f"pip install diffusers=={required_versions['diffusers']} transformers=={required_versions['transformers']} huggingface_hub=={required_versions['huggingface_hub']}")
        sys.exit(1)
except:
    pass


def run(prompt: str):
    print(f'开始生成图片....{prompt}')
    try:
        # 4. 加载文本编码器（保持高精度）
        text_encoder = ChineseCLIPTextModel.from_pretrained(
            model_path,
            subfolder="text_encoder",
            torch_dtype=torch.float32  # CPU保持float32，保证精度
        ).to(device)

        # 加载处理器
        processor = ChineseCLIPProcessor.from_pretrained(
            model_path,
            subfolder="tokenizer"
        )

        # 5. 加载Stable Diffusion管道（配置高质量采样器）
        # 自定义调度器：DDIM采样器，调整参数提升质量
        scheduler = DDIMScheduler.from_pretrained(
            model_path,
            subfolder="scheduler",
            beta_end=0.012,
            beta_schedule="scaled_linear",
            steps_offset=1  # 修复旧版本调度器的步数偏移问题
        )

        pipe = StableDiffusionPipeline.from_pretrained(
            model_path,
            text_encoder=text_encoder,
            scheduler=scheduler,  # 使用自定义调度器
            torch_dtype=torch.float32,
            device_map=None,
            low_cpu_mem_usage=True,
            ignore_mismatched_sizes=True,
            safety_checker=None
        ).to(device)

        # 6. 优化提示词处理（增强文本理解）
        def preprocess_prompt(prompt):
            inputs = processor(
                text=prompt,
                padding="max_length",
                max_length=77,  # 保留完整文本长度，提升语义理解
                truncation=True,
                return_tensors="pt"
            )
            return inputs.input_ids.to(device)

        # 7. 高质量生成参数配置
        # 正面提示词：增加细节描述（光影、材质、风格等）
        # prompt = (
        #     "明月松间照，清泉石上流。插画风格，"
        #     "高清细节，8K分辨率，柔和光影，层次感强，"
        #     "松树纹理清晰，泉水有透明质感，石头有自然纹理，"
        #     "古风色调，青绿色为主，构图平衡"
        # )
        # # 负面提示词：明确排除低质量特征
        # negative_prompt = (
        #     "模糊，低分辨率，变形，噪点，色偏，"
        #     "比例失调，细节丢失，笔触粗糙，"
        #     "不自然的光线，杂乱无章"
        #     "文字印章图案"
        # )

        # prompt = (
        #     "武松打虎，插画"
        # )
        negative_prompt = (
            "blood, gore, violence, murder, kill, dead, corpse, "
            "horrible, frightening, scary, monster, ghost, skeleton, zombie, "
            "sex, nudity, pornography, adult, erotic, mature, "
            "drugs, alcohol, smoking, tobacco, illegal, "
            "dark, night, storm, thunder, lightning, apocalypse, disaster, "
            "gun, knife, sword, bomb, explosion, firearm, "
            "mean, angry, sadistic, hostile, aggressive, bullying, "
            "dangerous, unsafe, hazardous, poison, toxic, pollution"
        )

        prompt_ids = preprocess_prompt(prompt)
        negative_ids = preprocess_prompt(negative_prompt)

        # 8. 高质量生成设置（核心优化）
        image = pipe(
            prompt_embeds=text_encoder(prompt_ids)[0],
            negative_prompt_embeds=text_encoder(negative_ids)[0],
            num_inference_steps=50,  # 增加步数（80步，细节更丰富）
            guidance_scale=8.0,  # 提高引导尺度（更严格遵循提示词）
            height=512,  # 标准高清分辨率（512x512，比256x256细节多4倍）
            width=512,
            eta=0.0,  # 采样噪声控制，降低随机性
            num_images_per_prompt=1  # 单次生成1张，集中资源
        ).images[0]

        # 保存高质量结果
        output_path = "high_quality_illustration.png"
        image.save(output_path)
        print(f"\n🎉 高质量图像生成成功，已保存为 {output_path}")

        # 正确的Base64转换流程：
        # 1. 创建内存缓冲区
        buffer = BytesIO()
        # 2. 将Image对象写入缓冲区（指定格式，如PNG/JPEG）
        image.save(buffer, format="PNG")  # 格式需与图像类型匹配
        # 3. 从缓冲区获取二进制数据
        image_bytes = buffer.getvalue()
        # 4. 编码为Base64字符串
        base64_str = base64.b64encode(image_bytes).decode("utf-8")
        return base64_str
    except Exception as e:
        print(f"\n运行时错误: {str(e)}")
        try:
            print(f"提示词ID形状: {prompt_ids.shape}")
            print(f"文本编码输出形状: {text_encoder(prompt_ids)[0].shape}")
        except:
            pass
        sys.exit(1)